import dashscope
import os
import base64
import io
import requests
from PIL import Image
from dotenv import load_dotenv
import time

# 加载环境变量
load_dotenv()

class NewsGenerator:
    def __init__(self, api_key=None):
        # 初始化API密钥
        self.api_key = api_key or os.getenv('DASHSCOPE_API_KEY')
        if not self.api_key:
            raise ValueError("请提供通义大模型的API密钥")
        
        # 配置API密钥
        dashscope.api_key = self.api_key
    
    def compress_image_for_vision(self, image: Image.Image, max_size=(1024, 1024)):
        """
        压缩图片以适应视觉模型输入限制
        - 缩放到最大 1024x1024
        - 转为 RGB（避免 RGBA 导致体积变大）
        - 使用 JPEG 格式减少体积
        """
        # 转为RGB(PNG 可能带透明通道，体积大)
        if image.mode in ("RGBA", "P"):
            image = image.convert("RGB")
        # 缩放
        image.thumbnail(max_size, Image.Resampling.LANCZOS)
        return image

    def analyze_image(self, image):
        """
        分析图片内容
        参数:
            image: PIL Image对象
        返回:
            str: 图片分析结果或"无法分析图片内容"
        """
        try:
            # 压缩图片
            compressed_img = self.compress_image_for_vision(image, max_size=(1024, 1024))

            # PIL Image转为字节流
            img_byte_arr = io.BytesIO()
            compressed_img.save(img_byte_arr, format='JPEG', quality=85)
            img_byte_arr.seek(0)
            
            # 转换为base64编码
            image_base64 = base64.b64encode(img_byte_arr.read()).decode('utf-8')
            
            # 调用qwen3-vl-plus视觉大模型分析图片信息
            response = dashscope.MultiModalConversation.call(
                model='qwen3-vl-plus',
                messages=[
                    {
                        'role': 'user',
                        'content': [
                            {'text': '请详细分析这张图片的内容，包括场景、人物、表情、动作、物品、环境、活动等所有关键信息。请提供专业、全面但简洁精炼的描述（控制在100字左右），重点提取用于校园新闻报道相关的关键辅助信息。'}, 
                            {'image': f"data:image/jpeg;base64,{image_base64}"}
                        ]
                    }
                ]
            )
            
            if response.status_code == 200:
                content = response.output.choices[0].message.content
                # 确保返回的是字符串而不是列表
                if isinstance(content, list) and len(content) > 0 and isinstance(content[0], dict) and 'text' in content[0]:
                    return content[0]['text']
                elif isinstance(content, str):
                    return content
                else:
                    return str(content)  # 兜底转换
            else:
                print(f"图片分析失败: {response.message}, 状态码: {response.status_code}")
                return "无法分析图片内容"
        except Exception as e:
            print(f"图片分析出错: {e}")
            return "无法分析图片内容"
            
    def generate_image(self, news_content, news_title, max_retries=3):
        """
        生成新闻AI配图
        参数:
            news_content: 新闻内容
            news_title: 新闻标题
            max_retries: 最大重试次数
        返回:
            PIL.Image对象或None(如果生成失败)
        """
        # 构建文生图提示词
        image_gen_prompt = f"为校园新闻稿《{news_title}》创建一张高质量的校园新闻风格配图。基于新闻内容：{news_content} 创建一张构图平衡的图片，风格鲜明、色彩协调、细节丰富，能准确传达新闻内容主题和校园氛围。"
        
        # 指数退避重试策略
        retry_delays = [2, 4, 8]  # 秒
        if max_retries > len(retry_delays):
            retry_delays.extend([retry_delays[-1]] * (max_retries - len(retry_delays)))
            
        for attempt in range(max_retries):
            if attempt > 0:
                delay = retry_delays[attempt-1] if attempt-1 < len(retry_delays) else retry_delays[-1]
                print(f"正在进行第 {attempt+1} 次尝试... (等待 {delay:.1f} 秒后重试)")
                time.sleep(delay)
            
            try:
                # 调用wan2.2-t2i-flash图片生成模型
                response = dashscope.ImageSynthesis.call(
                    model="wan2.2-t2i-flash",
                    prompt=image_gen_prompt,
                    n=1,
                    size='1280*720'
                )
                print(response)
                if response.status_code == 200 and response.output and response.output.results:
                    image_url = response.output.results[0].url
                    img_data = requests.get(image_url, timeout=10).content
                    return Image.open(io.BytesIO(img_data))
                else:
                    raise Exception(f"图片生成失败: {response.message if hasattr(response, 'message') else '未知错误'}")
            except Exception as e:
                print(f"图片生成出错: {str(e)}")
        
        # 所有重试都失败
        print(f"已完成所有 {max_retries} 次尝试，图片生成失败。")
        return None
    
    def generate_news(self, news_type, title_keywords, key_persons, time_location, content_points, language_style, max_length, additional_requirements, temperature=0.7, model="qwen-turbo", image_analyses=None):
        """
        生成校园新闻，利用图片分析信息增强新闻质量
        参数:
            news_type: 新闻类型
            title_keywords: 标题关键词
            key_persons: 主要人物
            time_location: 时间地点
            content_points: 核心内容要点
            language_style: 语言风格
            max_length: 最大长度
            additional_requirements: 附加要求
            temperature: 温度参数（影响创造性）
            model: 使用的模型
            image_analyses: 上传图片的分析信息
        返回:
            tuple: (生成的新闻内容, 提取的新闻标题)
        """
        # 处理图片分析结果,仅保留有效信息
        image_analysis_text = []
        if image_analyses:
            for i, analysis in enumerate(image_analyses):
                if isinstance(analysis, str) and analysis != "无法分析图片内容":
                    image_analysis_text.append(f"图片{i+1}中蕴含信息: {analysis}")  
        
        # 构建图片分析信息段落
        image_info_section = ""
        if image_analysis_text:
            image_info_section = "8. 图片信息分析：以下是上传图片的分析结果，请参考这些信息补充，自然融入到新闻内容中：" + "".join(image_analysis_text)
        
        # 增强版提示词，更好地利用图片分析信息
        prompt = f"""
        请帮我生成一篇{news_type}类型的校园新闻稿。
        
        基础信息：
        1. 标题关键词：{title_keywords}
        2. 主要人物：{key_persons}
        3. 时间地点：{time_location}
        4. 核心内容要点：{content_points}
        5. 语言风格：{language_style}
        6. 字数要求：约{max_length}字
        7. 附加要求：{additional_requirements}
        {image_info_section}
        
        写作要求：
        1. 生成一个吸引人的新闻标题
        2. 包含完整的新闻结构：标题、导语、正文
        3. 如果上述的基础信息中有图片信息分析，请将图片信息分析数据作为信息补充，丰富生成的新闻内容
        4. 确保内容连贯、专业、符合新闻写作规范
        5. 请严格按照以下格式输出：
           - 第一行必须是标题，并且必须以'# '开头（#号后面跟一个空格）
           - 标题后面空一行
           - 然后是新闻正文内容
        6. 文章内容要详实、生动，充分利用提供的所有信息
        
        请生成完整的新闻稿，不要添加任何额外说明。
        """

        # 调用通义大模型
        response = dashscope.Generation.call(
            model=model,
            prompt=prompt,
            max_tokens=max_length * 2,  # 预留足够空间
            temperature=temperature,
            top_p=0.9
        )
        
        # 检查响应
        if response.status_code == 200:
            news_text = response.output.text
            # 提取新闻标题
            news_lines = news_text.strip().split('\n')
            news_title = ""
            title_found_index = -1
            
            # 优先查找以#开头的标题行
            for i, line in enumerate(news_lines):
                leading_removed = line.lstrip()
                if leading_removed.startswith('#'):
                    # 处理标题内容
                    if len(leading_removed) > 1 and leading_removed[1] == ' ':
                        news_title = leading_removed[2:].strip()
                    else:
                        news_title = leading_removed[1:].strip()
                    title_found_index = i
                    break
            
            # 找不到标题：使用标题关键词作为默认标题
            if not news_title:
                news_title = title_keywords
            
            # 移除标题部分，只保留正文
            if title_found_index >= 0:
                # 移除标题行和可能的空行
                remaining_lines = []
                skip_next_blank = True  # 是否跳过标题后的空行
                
                for i, line in enumerate(news_lines):
                    if i > title_found_index:  # 只处理标题之后的行
                        if skip_next_blank and not line.strip():
                            continue
                        skip_next_blank = False  # 之后的内容都保留
                        remaining_lines.append(line)
                
                # 重建不包含标题的新闻正文
                news_text = '\n'.join(remaining_lines)
                
            return news_text, news_title
        else:
            raise Exception(f"生成失败: {response.message}")